package com.cqs.leetcode.search.binary;

import java.util.*;

/**
 * lixw
 * 2023/10/13
 */
public class Solution373 {

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        int l = nums1[0] + nums2[0], r = nums1[nums1.length - 1] + nums2[nums2.length - 1];
        while (l < r) {
            int mid = (r - l) / 2 + l, check = check(nums1, nums2, mid, k);
            if (check < 0) {
                l = mid + 1;
            } else {
                r = mid;
            }
        }
        List<List<Integer>> result = new ArrayList<>();
        for (int i = 0; i < nums1.length; i++) {
            for (int j = 0; j < nums2.length && nums1[i] + nums2[j] <= r; j++) {
                List<Integer> list = new ArrayList<>();
                list.add(nums1[i]);
                list.add(nums2[j]);
                result.add(list);
            }
        }
        Collections.sort(result, Comparator.comparingInt(o -> o.get(0) + o.get(1)));
        return k >= result.size() ? result : result.subList(0, k);
    }

    /**
     * 是否存在k个数据对元素<=target
     *
     * @param nums1
     * @param nums2
     * @param target
     * @param k
     * @return
     */
    private int check(int[] nums1, int[] nums2, int target, int k) {
        int cnt = 0;
        for (int i = nums1.length - 1; i >= 0; --i) {
            //从num2中找到最后右边一个 使得num1[i] + num2[x] <= target;
            int l = 0, r = nums2.length - 1;
            while (l < r) {
                int mid = (r + l + 1) / 2;
                if (nums1[i] + nums2[mid] > target) {
                    r = mid - 1;
                } else {
                    l = mid;
                }
            }
            l = nums1[i] + nums2[l] <= target ? l : -1;
            cnt += l + 1;
            if (cnt >= k) return 1;
        }
        return -1;
    }
}
